package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"sync"
	"time"
)

func main() {
	fmt.Println("Hello, Here is TCP Scanner")
	hostname := flag.String("hostname", "", "hostname to test")
	startPort := flag.Int("start-port", 80, "the port on which yje scanning starts")
	endPort := flag.Int("end-port", 100, "the port form which the scanning ends")
	timeout := flag.Duration("timeout", time.Millisecond*1000, "timeout")
	// 自定义帮助信息
	flag.Usage = func() {
		_, _ = fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		_, _ = fmt.Fprintln(os.Stderr, "A simple TCP port scanner")
		_, _ = fmt.Fprintln(os.Stderr, "\nOptions:")
		flag.PrintDefaults()
		_, _ = fmt.Fprintln(os.Stderr, "\nExample:")
		_, _ = fmt.Fprintln(os.Stderr, "  ", os.Args[0], "-hostname example.com -start-port 80 -end-port 100")
		_, _ = fmt.Fprintln(os.Stderr, "  ", os.Args[0], "-hostname localhost -start-port 1 -end-port 1024 -timeout 100ms")
	}

	flag.Parse()
	// 检查是否需要显示帮助
	if *hostname == "" {
		flag.Usage()
		os.Exit(1)
	}
	fmt.Printf("hostname: %v, ports[%v~%v]\n", *hostname, *startPort, *endPort)

	var ports []int
	mutex := &sync.Mutex{}
	wg := &sync.WaitGroup{}
	for port := *startPort; port <= *endPort; port++ {
		wg.Add(1)
		go func(p int) {
			opened := isOpen(*hostname, p, *timeout)
			if opened {
				mutex.Lock()
				ports = append(ports, p)
				mutex.Unlock()
			}
			wg.Done()
		}(port)
	}

	wg.Wait()
	fmt.Println("opened ports: ", ports)
}

func isOpen(host string, port int, timeout time.Duration) bool {
	time.Sleep(time.Microsecond * 1)
	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", host, port), timeout)
	if err == nil {
		_ = conn.Close()
		return true
	}
	return false
}
